{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "14fd8270-c7eb-4b5b-aa4e-1168d3fb20b4",
   "metadata": {},
   "source": [
    "## RLAIF Example - Finetuning with Sagemaker\n",
    "\n",
    "This notebook demonstrates basic user flow for RLAIF Finetuning from a model available in Sagemaker Jumpstart.\n",
    "Information on available models on jumpstart: https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-foundation-models-latest.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0be8f0fd-79f6-49ec-8658-edb12a5e86fd",
   "metadata": {},
   "source": [
    "### Setup and Configuration\n",
    "\n",
    "Initialize the environment by importing necessary libraries and configuring AWS credentials"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ],
   "id": "4a05468e7078023e",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "cec1af2d-c0c1-4348-8ee7-502a6d7ee2d0",
   "metadata": {},
   "source": [
    "#!/usr/bin/env python3\n",
    "\n",
    "from sagemaker.train.rlaif_trainer import RLAIFTrainer\n",
    "from sagemaker.train.configs import InputData\n",
    "from rich import print as rprint\n",
    "from rich.pretty import pprint\n",
    "from sagemaker.core.resources import ModelPackage\n",
    "import os\n",
    "#os.environ['SAGEMAKER_REGION'] = 'us-east-1'\n",
    "#os.environ['SAGEMAKER_STAGE'] = 'prod'\n",
    "\n",
    "import boto3\n",
    "from sagemaker.core.helper.session_helper import Session\n",
    "\n",
    "# For MLFlow native metrics in Trainer wait, run below line with approriate region\n",
    "os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-west-2.app.aws\"\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a835d03d-73a7-439d-a17a-7491df38872a",
   "metadata": {},
   "source": [
    "### Create RLAIFTrainer\n",
    "**Required Parameters** \n",
    "\n",
    "* `model`: base_model id on Sagemaker Hubcontent that is available to finetune (or) ModelPackage artifacts\n",
    "\n",
    "**Optional Parameters**\n",
    "* `reward_model_id`: Bedrock model id, supported evaluation models: https://docs.aws.amazon.com/bedrock/latest/userguide/evaluation-judge.html\n",
    "*  `reward_prompt`: Reward prompt ARN or builtin prompts refer: https://docs.aws.amazon.com/bedrock/latest/userguide/model-evaluation-metrics.html\n",
    "* `model_package_group_name`: ModelPackage group name or ModelPackageGroup\n",
    "* `mlflow_resource_arn`: MLFlow app ARN to track the training job\n",
    "* `mlflow_experiment_name`: MLFlow app experiment name(str)\n",
    "* `mlflow_run_name`: MLFlow app run name(str)\n",
    "* `training_dataset`: Training Dataset - either Dataset ARN or S3 Path of the dataset (Please note these are required for a training job to run, can be either provided via Trainer or .train())\n",
    "* `validation_dataset`: Validation Dataset - either Dataset ARN or S3 Path of the dataset\n",
    "* `s3_output_path`: S3 path for the trained model artifacts"
   ]
  },
  {
   "cell_type": "code",
   "id": "07aefa46-29f2-4fcf-86da-b0bd471e0a6a",
   "metadata": {},
   "source": [
    "# For fine-tuning \n",
    "rlaif_trainer = RLAIFTrainer(\n",
    "    model=\"meta-textgeneration-llama-3-2-1b-instruct\", # Union[str, ModelPackage] \n",
    "    model_package_group_name=\"sdk-test-finetuned-models\", # Make it Optional\n",
    "    reward_model_id='anthropic.claude-3-5-sonnet-20240620-v1:0',\n",
    "    reward_prompt='Builtin.Correctness',\n",
    "    #mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-rlaif-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-rlaif-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl\", #Optional[]\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing/output/\",\n",
    "    accept_eula=True\n",
    "    #sagemaker_session=sagemaker_session,\n",
    "    #role=\"arn:aws:iam::<>:role/Admin\"\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "8db04c11-fa5c-4716-be19-50e594e35811",
   "metadata": {},
   "source": [
    "### Discover and update Finetuning options\n",
    "\n",
    "Each of the technique and model has overridable hyperparameters that can be finetuned by the user."
   ]
  },
  {
   "cell_type": "code",
   "id": "b31d57c0-9777-428d-8792-557f7be4cfda",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "print(\"Default Finetuning Options:\")\n",
    "pprint(rlaif_trainer.hyperparameters.to_dict()) # rename as hyperparameters\n",
    "\n",
    "#set options\n",
    "rlaif_trainer.hyperparameters.get_info()\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "5f416275-65d4-4dfe-bb64-e17c5f34146c",
   "metadata": {},
   "source": [
    "#### Start RLAIF training\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "5d5fa362-0caf-412d-977c-5e47f0548ea5",
   "metadata": {},
   "source": [
    "\n",
    "training_job = rlaif_trainer.train(wait=True)\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "a0781a22-d9ea-4c9b-a854-5d7efde3539d",
   "metadata": {},
   "source": [
    "training_job = rlaif_trainer.train(wait=True)\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c34b93c8-2e4c-437a-8efb-b8475fb941f3",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "import re\n",
    "from sagemaker.core.utils.utils import Unassigned\n",
    "import json\n",
    "\n",
    "\n",
    "def pretty_print(obj):\n",
    "    def parse_unassigned(item):\n",
    "        if isinstance(item, Unassigned):\n",
    "            return None\n",
    "        if isinstance(item, dict):\n",
    "            return {k: parse_unassigned(v) for k, v in item.items() if parse_unassigned(v) is not None}\n",
    "        if isinstance(item, list):\n",
    "            return [parse_unassigned(x) for x in item if parse_unassigned(x) is not None]\n",
    "        if isinstance(item, str) and \"Unassigned object\" in item:\n",
    "            pairs = re.findall(r\"(\\w+)=([^<][^=]*?)(?=\\s+\\w+=|$)\", item)\n",
    "            result = {k: v.strip(\"'\\\"\") for k, v in pairs}\n",
    "            return result if result else None\n",
    "        return item\n",
    "\n",
    "    cleaned = parse_unassigned(obj.__dict__ if hasattr(obj, '__dict__') else obj)\n",
    "    print(json.dumps(cleaned, indent=2, default=str))\n",
    "\n",
    "pretty_print(training_job)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "3ad71843-ab14-43ce-9eea-5cd2fa46ab6d",
   "metadata": {},
   "source": [
    "### View any Training job details\n",
    "\n",
    "We can get any training job details and its status with TrainingJob.get(...)"
   ]
  },
  {
   "cell_type": "code",
   "id": "860fcbd0-d340-4bde-bfbc-224b4b8b0aed",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "from sagemaker.core.resources import TrainingJob\n",
    "\n",
    "response = TrainingJob.get(training_job_name=\"meta-textgeneration-llama-3-2-1b-instruct-rlaif-20251124140754\")\n",
    "pretty_print(response)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "94faaefd-6987-495b-b96a-e6ba08afb9ed",
   "metadata": {},
   "source": [
    "### Test RLAIF with Custom Reward Prompt\n",
    "\n",
    "Here we are providing a user-defined reward prompt/evaluator ARN"
   ]
  },
  {
   "cell_type": "code",
   "id": "05015ab9-3e70-4b61-affe-cd84ed4eccae",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "# For fine-tuning \n",
    "rlaif_trainer = RLAIFTrainer(\n",
    "    model=\"meta-textgeneration-llama-3-2-1b-instruct\", # Union[str, ModelPackage] \n",
    "    model_package_group_name=\"sdk-test-finetuned-models\", # Make it Optional\n",
    "    reward_model_id='anthropic.claude-3-5-sonnet-20240620-v1:0',\n",
    "    reward_prompt=\"arn:aws:sagemaker:us-west-2:<>:hub-content/sdktest/JsonDoc/rlaif-test-prompt/0.0.1\",\n",
    "    #mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-rlaif-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-rlaif-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl\", #Optional[str]\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing/output/\",\n",
    "    accept_eula=True\n",
    "    #sagemaker_session=sagemaker_session,\n",
    "    #role=\"arn:aws:iam::<>:role/service-role/AmazonSageMaker-ExecutionRole-20250731T162975\"\n",
    "    #role=\"arn:aws:iam::<>:role/Admin\"\n",
    ")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "029aa3cf-8a98-487b-8e21-445af9a72e91",
   "metadata": {},
   "source": [
    "training_job = rlaif_trainer.train(wait=True,\n",
    "                                   logs=True)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "6815868d-0490-43a4-9765-148c4b2ef4af",
   "metadata": {},
   "source": [
    "training_job = rlaif_trainer.train(wait=True)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "fc3e1c67-d9c5-429a-aed7-07b26106ef2e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "scrolled": true
   },
   "source": "pretty_print(training_job)",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
